from __future__ import print_function
import os
import numpy as np
import keras
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from keras.layers import Input
from keras.layers import Dense
from keras.layers import Dropout
from keras.applications.inception_v3 import InceptionV3
from keras.preprocessing import image
from keras.models import Model
from keras.preprocessing.image import ImageDataGenerator
from keras import backend as K
from keras.callbacks import EarlyStopping  
import argparse
#设置输入图片大小，InceptionV3的默认输入是（299，,299）
img_width, img_height = 299, 299
input_tensor = Input(shape=(img_width, img_height, 3))
#数据集路径
train_data_dir = os.path.join(os.getcwd(), 'dataset/train')
validation_data_dir = os.path.join(os.getcwd(), 'dataset/test')
#数据集大小
nb_train_samples = 155
nb_validation_samples =25
#训练的参数
num_class = 2
epochs = 30
batch_size = 16
#构建基础模型
base_model = InceptionV3(include_top = False, weights = 'imagenet',  input_tensor = input_tensor, pooling = 'avg')
#可视化层名称和层指数
for i, layer in enumerate(base_model.layers):
    print(i, layer.name)
#增加新的输出层
x = base_model.output
x = Dense(2048, activation='relu')(x)
x = Dropout(0.5)(x)
pre_out = Dense(num_class, activation='softmax')(x)
#reconstruct the model  重建模型
train_model = Model(base_model.input, outputs= pre_out, name='train_model')
#冻结前面层，对其余4层进行更新
for layer in base_model.layers[:4]:
    layer.trainable = False
for layer in base_model.layers[4:]:
    layer.trainable = True
#微调学习率
sgd = keras.optimizers.SGD(lr=0.0001, decay=1e-6, momentum=0.9, nesterov=True)
#重新编译模型
train_model.compile(loss='categorical_crossentropy', 
                   optimizer=sgd, metrics=['accuracy'])
train_model.summary()
#图像数据扩充
#训练数据生成器
train_datagen = ImageDataGenerator(rotation_range=60,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True,
                                   vertical_flip=True,
                                   rescale=1. / 255,
                                   fill_mode='nearest')

#测试数据生成器
test_datagen = ImageDataGenerator(rescale=1. / 255)

train_generator = train_datagen.flow_from_directory(
    'dataset/train',
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode='categorical')

validation_generator = test_datagen.flow_from_directory(
    validation_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode='categorical')


#设置early-stopping
early_stopping = EarlyStopping(monitor='val_loss', patience=3)

#在训练数据集上训练模型
hist = train_model.fit_generator(
    train_generator,
    steps_per_epoch=nb_train_samples //batch_size,
    epochs=epochs,
    validation_data=validation_generator,
    validation_steps=nb_validation_samples //batch_size, 
    callbacks=[early_stopping])

#打印准确率并存储到acc.txt中
f = open('acc.txt','w')
f.write(str(hist.history['accuracy']))
f.close()
#print val_acc and stored into val_acc.txt
f = open('val_acc.txt','w')
f.write(str(hist.history['val_accuracy']))
f.close()
#print val_loss and stored into val_loss.txt   
f = open('val_loss.txt', 'w')
f.write(str(hist.history['val_loss']))
f.close()

# 评估模型
evaluation = train_model.evaluate_generator(validation_generator,
                                      steps=nb_validation_samples //batch_size)

print('Model Accuracy = %.4f' % (evaluation[1]))

#保存模型
save_dir = os.path.join(os.getcwd(), 'googlenet_model')
model_name = 'keras_googlenet_trained_model.h5'
if not os.path.isdir(save_dir):
    os.makedirs(save_dir)
model_path = os.path.join(save_dir, model_name)
train_model.save(model_path)
print('save trained model at %s' %model_path)
# 绘制训练损失和准确性
#N = args["epochs"]
#plt.style.use("ggplot")
plt.figure()
#plt.plot(np.arange(0, N), trainingModel.history["loss"], label="train_loss")
#plt.plot(np.arange(0, N),trainingModel.history["val_loss"], label="val_loss")
plt.plot(np.arange(0, epochs), hist.history["accuracy"], label="train_acc")
plt.plot(np.arange(0, epochs), hist.history["val_accuracy"], label="val_acc")
plt.title("Training Loss and Accuracy on Dataset")
plt.xlabel("Epoch ")
plt.ylabel("Accuracy")
plt.legend(loc="lower left")
plt.savefig(args["plot"])
 